import glob
import traceback
from pathlib import Path

import shapefile

from .excel import Excel


class ToShp:
    def __init__(
            self,
            filedir: str,
            saveFile: str,
            columns: list,
            *,
            encoding='UTF-8',
            _type='.xlsx',
            name: str = '名称',
            headers=None,
            **kwargs
    ):
        self.filedir = filedir
        self.saveFile = saveFile
        self._type = _type
        self.columns = columns
        self.name = name
        self.headers = headers

        self.cycle_number = 0

        self.shpWrite = shapefile.Writer(saveFile, encoding=encoding, **kwargs)
        self.xlsxList = self.getXlsxFileList()

    def setSHP(self, file, element_type):
        excel = Excel(file)
        if self.cycle_number == 0:
            if self.headers is None:
                header = excel.getHeaders()
            else:
                header = self.headers
            header.append(self.name)
            for title in header:
                self.shpWrite.field(title)
        dataList = excel.to_dict()
        name = Path(file).name.split('.')[0]
        self.shpWrite.record(*list(dataList[0].values()), name)
        getattr(self.shpWrite, element_type)([excel.getColumnsData(self.columns)])
        # self.shpWrite.line([excel.getColumnsData(self.columns)])
        self.cycle_number += 1

    def getXlsxFileList(self):
        xlsxFileList = []

        def getXlsxFile(file):
            for i in glob.glob(file + '/*'):
                if Path(i).is_dir():
                    getXlsxFile(i)
                else:
                    if Path(i).name.endswith(self._type):
                        xlsxFileList.append(i)

        getXlsxFile(self.filedir)
        return xlsxFileList

    def run(self, element_type: str = 'line'):
        """
        开始
        :param element_type: 要素类型
             point: 点, (x, y)
             pointm: 点,(x, y, m)
             pointz: 点, (x, y, z, m)
             line: 线, (x, y)
             linem: 线,(x, y, m)
             linez: 线, (x, y, z, m)
             poly: 面, (x, y)
             polym: 面,(x, y, m)
             polyz: 面, (x, y, z, m)
        :return:
        """
        for index, xlsx in enumerate(self.xlsxList):
            index = index + 1
            try:
                self.setSHP(xlsx, element_type)
                yield xlsx, '成功', int(index / len(self.xlsxList) * 100)
            except Exception as e:
                error = traceback.format_exc()
                yield xlsx, error, int(index / len(self.xlsxList) * 100)
                continue

    def close(self):
        self.shpWrite.close()

    def __enter__(self):
        self.close()


if __name__ == '__main__':
    shp = ToShp(r'F:\PythonObject\pythonProject1\线路轨迹', r'F:\PythonObject\pythonProject1\test\test.shp', ['经度', '纬度'])
    print(shp.run())
